import matplotlib.pyplot as plt
import math
import random
import numpy as np

from data_process import getUbAndLb

"""
函数里面所有以plot开头的函数都可以注释掉，没有影响
求解的目标表达式为：
y = 10 * math.sin(5 * x) + 7 * math.cos(4 * x)  x belongs to (0,10)

需要将X拓展为三维坐标，并且重写定义目标函数
"""


class SA():
    def __init__(self, pN, dim, max_iter, ub, lb, T_init=100, T_min=1e-3, alpha=0.90):
        '''
        :param pN: 粒子数量
        :param dim:  粒子的维度（3维）
        :param max_iter: 最大迭代数量
        :param ub: 取值上限
        :param lb: 取值下限
        :param T_init: 初始最大温度
        :param T_min: 最小温度，退出循环条件
        :param alpha: 降维系数
        '''
        self.T_init = T_init
        self.alpha = alpha
        self.T_min = T_min
        self.results = []  # 存x，y
        self.pN = pN
        self.dim = dim
        self.max_iter = max_iter
        assert len(lb) == len(ub), 'Lower- and upper-bounds must be the same length'
        assert np.all(ub > lb), 'All upper-bound values must be greater than lower-bound values'
        self.ub = np.array(ub)
        self.lb = np.array(lb)
        self.X = self.lb + np.random.rand(self.pN, self.dim) * (self.ub - self.lb)
        self.X_best = np.zeros(self.dim)
        self.y = np.zeros(self.pN)
        self.y_best = 0
        self.records = []

        # plot_obj_func()

    def get_delta_X(self):
        return np.random.rand(self.dim) - 0.5

    def objectFunc(self, X):
        return self.getSutton_Chen(X)

    def getSutton_Chen(self, location):
        c = 2
        a = 1
        m = 1
        sutton = 0
        p_i = 0
        p_1 = 0
        for particle in self.X:
            if (particle == location).all():
                continue
            r = np.linalg.norm(particle - location)
            p_i += (a / r) ** m
            p_1 += (a / r)
        sutton += 4 * (0.5 * p_1 - c * np.sqrt(p_i))
        return sutton

    def run(self):
        while self.T_init > self.T_min:
            # 每个温度迭代50次，找最优解
            for t in range(self.max_iter):
                for i in range(self.pN):
                    self.y[i] = self.objectFunc(self.X[i])
                    delta_X = self.get_delta_X()  # 自变量进行波动
                    # 自变量变化后仍要求在[0,10]之间
                    X_new = np.zeros_like(self.X[i])
                    for d in range(self.dim):
                        if self.lb[d] < (self.X[i][d] + delta_X[d]) < self.ub[d]:
                            X_new[d] = self.X[i][d] + delta_X[d]
                        else:
                            X_new[d] = self.X[i][d] - delta_X[d]
                    # 目标函数
                    y_new = self.objectFunc(X_new)
                    # 以上为找最大值，要找最小值就把>号变成<
                    if y_new < self.y[i] or math.exp(-(y_new - self.y[i]) / self.T_init) > random.random():
                        # 有新值被接受
                        self.X[i] = X_new
                        self.y[i] = y_new
                        if y_new < self.y_best:
                            self.X_best = self.X[i]
                            self.y_best = self.y[i]
                print("------------当前位置----------")
                print(self.X[0], end=" ")
                print(end="\n")
                print("------------当前最优值---------")
                print(self.y_best)
                self.records.append(self.y_best)
            self.results.append((self.X, self.y))
            self.T_init *= self.alpha

        print('局部最优解 x:{},y:{}'.format(self.results[-1][0], self.results[-1][1]))
        print('全局最优解 x:{},y:{}'.format(self.X_best, self.y_best))
        # plot_final_result(self.results)
        plot_iter_curve(self.records)


# 看看我们要处理的目标函数
def plot_obj_func():
    """y = 10 * math.sin(5 * x) + 7 * math.cos(4 * x)"""
    X1 = [i / float(10) for i in range(0, 100, 1)]
    Y1 = [10 * math.sin(5 * x) + 7 * math.cos(4 * x) for x in X1]
    plt.title("object func")
    plt.plot(X1, Y1)
    plt.show()


# 看看最终的迭代变化曲线
def plot_iter_curve(results):
    X = [i for i in range(len(results))]
    Y = results
    plt.plot(X, Y)
    plt.title("the curve of target value")
    plt.show()


def plot_final_result(results):
    X1 = [i / float(10) for i in range(0, 100, 1)]
    Y1 = [10 * math.sin(5 * x) + 7 * math.cos(4 * x) for x in X1]
    plt.plot(X1, Y1)
    plt.title("the curve of X versus Y")
    plt.scatter(results[-1][0], results[-1][1], c='r', s=10)
    plt.show()


if __name__ == '__main__':
    # for i in range(100):
    # main()
    ub, lb = getUbAndLb()
    sa = SA(20, 3, 20, ub=ub, lb=lb)
    sa.run()
